package learning.crf.inference;

import learning.crf.model.Model;

public class Viterbi {

	private IParseScorer parseScorer;
	
	private Model model;
	
	public Viterbi(Model model, IParseScorer parseScorer) {
		this.model = model;
		this.parseScorer = parseScorer;
	}
	
	public Parse parse(int length) {
	
		int[][] parses = new int[length][model.numStates];
		float[][] scores = new float[length][model.numStates];
		
		for (int s = 0; s < model.numStates; s++)
			scores[0][s] = parseScorer.scoreTransition(0, model.START_STATE, s);
		
		for (int time = 1; time < length; time++) {
			for (int s = 0; s < model.numStates; s++) {
				// determine max scoring previous state
				int maxPrev = 0;
				float maxScore = Float.NEGATIVE_INFINITY;
				
				for (int ps = 0; ps < model.numStates; ps++) {
					float score = scores[time-1][ps] + 
						parseScorer.scoreTransition(time, ps, s);
					if (score > maxScore) {
						maxScore = score;
						maxPrev = ps;
					}			
				}
				scores[time][s] = maxScore;
				parses[time][s] = maxPrev;
			}
		}
		
		// determine highest scoring end state
		int maxS = 0;
		for (int s = 0; s < model.numStates; s++)
			if (scores[length-1][s] > scores[length-1][maxS])
				maxS = s;
		
		// write best state sequence
		int[] states = new int[length];
		states[length-1] = maxS;
		for (int j=length-2; j >= 0; j--)
			states[j] = parses[j+1][states[j+1]];
		
		return new Parse(states, scores[length-1][maxS]);
	}
	
	
	
}
